from qiskit import QuantumCircuit
def create_ghz(drawing=False):
    """ Generate a QuantumCircuit for a 3 qubit GHZ State and measure it. If `drawing` is True, return both the circuit object and the Matplotlib drawing of the circuit, otherwise return just the circuit object.
    """
    ghz = QuantumCircuit(3)
    ghz.h(0)
    ghz.cx(0, 1)
    ghz.cx(0, 2)
    ghz.measure_all()
    if drawing:
        return ghz, ghz.draw(output="mpl")
    return ghz

def check(candidate):
    from qiskit.quantum_info import Statevector
    import math
    import matplotlib

    def check_circuit(circuit):
        assert circuit.data[-1].operation.name == "measure"
        circuit.remove_final_measurements()
        ghz_statevector = (
            Statevector.from_label("000") + Statevector.from_label("111")
        ) / math.sqrt(2)
        assert Statevector.from_instruction(circuit).equiv(ghz_statevector)

    circuit = candidate()
    check_circuit(circuit)

    circuit, drawing = candidate(drawing=True)
    check_circuit(circuit)
    assert isinstance(drawing, matplotlib.figure.Figure)

    print("OK")

if __name__ == '__main__':
    check(create_ghz)